--- The strict 'State' 'Monad'

module frege.control.monad.State
         inline  (
                 State.State.>>=, 
                 State.State.>>,
                 State.State.pure, 
                 State.State.get,
                 -- State.run,
                 StateT.>>=, 
                 StateT.>>, 
                 StateT.pure, StateT.get,
                 StateT.lift, StateT.liftIO) 
    where

import frege.control.monad.trans.MonadTrans
import frege.control.monad.trans.MonadIO
-- class MonadState m where
--     get :: m s s
--     put :: s -> m s ()
--     modify :: (s -> s) -> m s () 


{--
    @State s a@ is an abstract data type that resembles a stateful computation
    with State _s_ and result _a_,
    i.e. functions of type @s -> (a, s)@
    where the State is immutable.
 -}
abstract data State s a = State { !fun :: s -> (a,s) } where
    --- run a stateful computation
    public run (State x) s = x s

    --- return the state from the internals of the monad
    public get   = State (\s -> (s, s))
    
    --- replace the 'State' inside the monad 
    public put x = State (\s -> ((), x))
    --- modify the 'State'
    public modify f = State (\state -> ((), f state))
    --- lift a value to the 'State' monad
    public pure !a =  State (\s -> (a, s))
    --- monadic bind for the 'State' monad
    public State fun >>= k  =  State (\s -> case fun s of 
                                        (v, s') -> case k v of State q -> q s'
                                    )
    public a >> b = a  >>=  \_ → b


instance Monad (State s)

--- make an instance of 'State' using an appropriate function, for example
--- > state random
state f = State.State f 

runState    = State.run
evalState s = fst . State.run s
execState s = snd . State.run s


runStateT :: Monad m ⇒ StateT s m v → s → m (v, s)
runStateT   = StateT.run

evalStateT ∷ Monad m ⇒ StateT s m v → s → m v
evalStateT s z = fst <$> runStateT s z

execStateT ∷ Monad m ⇒ StateT s m v → s → m s
execStateT s z = snd <$> runStateT s z
-- state       = State.State

--- > StateT s m a
--- A stateful computation that transforms an inner 'Monad' _m_
abstract data StateT s m a = StateT { !run ::  s -> m (a,s) } where
    public get   = StateT (\s -> pure (s, s)) 
    public put s = StateT (\_ -> pure ((), s))
    public modify f = StateT (\s -> pure ((), f s))



instance (Monad m) => Monad (StateT s m) where
    a >> b = a  >>=  \_ → b
    pure !a = StateT.StateT (\s -> pure (a,s))
    (>>=) :: Monad m ⇒ StateT s m a → (a→StateT s m b) → StateT s m b
    (StateT.StateT x) >>= f = StateT.StateT (\s -> do
        (v,s') <- x s          -- get new value and state
        case f v of
            StateT.StateT y ->  y s'     -- pass them to f
      )

instance (MonadPlus m) => MonadPlus (StateT s m) where
    mzero = StateT.StateT (\s -> mzero)
    (StateT.StateT x1) `mplus` (StateT.StateT x2) = StateT.StateT (\s -> (x1 s) `mplus` (x2 s))

instance (MonadAlt m) => MonadAlt (StateT s m) where
    pzero = StateT.StateT (\s -> pzero)
    (StateT.StateT x1) <|> (StateT.StateT x2) = StateT.StateT (\s -> (x1 s) <|> (x2 s))
    (StateT.StateT x1) <+> (StateT.StateT x2) = StateT.StateT (\s -> (x1 s) <+> (x2 s))

instance MonadTrans (StateT s) where
    lift c = StateT.StateT (\s -> c >>= (\x -> pure (x,s)))

instance (MonadIO m) =>  MonadIO (StateT s m) where
    liftIO = lift . liftIO

promote st = StateT.StateT (\s -> pure (State.run st s))